"""
仿真模块 - 定义和管理仿真环境及其规则

包含:
- Environment: 仿真环境类，定义仿真区域和边界规则
- SimulationEngine: 仿真引擎类，管理实体和更新仿真状态
"""

import numpy as np
from typing import List, Optional, Dict, Any, Tuple
from .entities import USV, Target, Entity


class Environment:
    """
    仿真环境类，表示USV和目标运行的水面环境
    
    属性:
        width: 环境宽度
        height: 环境高度
        boundary_mode: 边界处理模式 ('wrap', 'bounce', 'open')
    """
    
    def __init__(self, width: float = 100.0, height: float = 100.0, 
                 boundary_mode: str = 'wrap'):
        """
        初始化仿真环境
        
        :param width: 环境宽度
        :param height: 环境高度
        :param boundary_mode: 边界处理模式，可选 'wrap'(环绕), 'bounce'(反弹), 'open'(开放)
        """
        self.width = width
        self.height = height
        self.boundary_mode = boundary_mode
        
    def apply_boundary_rules(self, entity: Entity):
        """
        对实体应用边界规则
        
        :param entity: 需要应用边界规则的实体
        """
        x, y = entity.position
        
        if self.boundary_mode == 'wrap':  # 环绕模式
            entity.position[0] = x % self.width
            entity.position[1] = y % self.height
            
        elif self.boundary_mode == 'bounce':  # 反弹
            if x < 0:
                entity.position[0] = 0
                entity.velocity[0] *= -1
            elif x > self.width:
                entity.position[0] = self.width
                entity.velocity[0] *= -1
                
            if y < 0:
                entity.position[1] = 0
                entity.velocity[1] *= -1
            elif y > self.height:
                entity.position[1] = self.height
                entity.velocity[1] *= -1
                
        # 'open' 模式不做处理


class SimulationEngine:
    """
    仿真引擎类，管理整个仿真系统
    
    属性:
        env: 仿真环境
        usvs: 无人艇列表
        targets: 目标列表
        time: 当前仿真时间
    """
    
    def __init__(self, env: Environment):
        """
        初始化仿真引擎
        
        :param env: 仿真环境实例
        """
        self.env = env
        self.usvs: List[USV] = []
        self.targets: List[Target] = []
        self.time = 0.0
        self.metrics: Dict[str, List[float]] = {
            'avg_distance_to_target': [],
            'target_surrounded_ratio': [],
            'usv_clustering': []
        }
        
        # 目标分配数据
        self.usv_assignments = {}
        self.usv_leadership = {}
        self.encirclement_sectors={}
        self.reassignment_cooldown = 0.0
        self.enable_reassignment=True
        
    def add_usv(self, usv: USV):
        """
        添加无人艇到仿真系统
        
        :param usv: USV实例
        """
        self.usvs.append(usv)
        
    def add_target(self, target: Target):
        """
        添加目标到仿真系统
        
        :param target: Target实例
        """
        self.targets.append(target)
        
    def add_usv_fleet(self, count: int, spawn_area: Tuple[float, float, float, float],
                      speed_range: Tuple[float, float] = (0.5, 1.5),
                      perception_radius: float = 5.0):
        """
        批量添加一组无人艇
        
        :param count: 添加的无人艇数量
        :param spawn_area: 生成区域 (x_min, y_min, x_max, y_max)
        :param speed_range: 初始速度范围 (min_speed, max_speed)
        :param perception_radius: 感知半径
        """
        x_min, y_min, x_max, y_max = spawn_area
        for _ in range(count):
            # 随机位置
            pos = np.array([
                np.random.uniform(x_min, x_max),
                np.random.uniform(y_min, y_max)
            ])
            
            # 随机初始速度
            speed = np.random.uniform(speed_range[0], speed_range[1])
            angle = np.random.uniform(0, 2 * np.pi)
            vel = np.array([
                speed * np.cos(angle),
                speed * np.sin(angle)
            ])
            
            # 创建并添加USV
            usv=USV(position=pos, velocity=vel, perception_radius=perception_radius)
            self.add_usv(usv)
            
    def update(self, dt: float):
        """
        更新整个仿真状态
        
        :param dt: 时间步长，控制仿真进度的增量
        :return: None
        """
        # 重新分配检查
        self.reassignment_cooldown -= dt
        if self.reassignment_cooldown <= 0:
            self._assign_usvs_to_targets()
            self.reassignment_cooldown = 5.0  # 5秒冷却
        
        # 更新目标
        for target in self.targets:
            # 创建目标特定的追捕艇列表
            pursuing_usvs = [usv for usv in self.usvs 
                            if self.usv_assignments.get(id(usv)) == id(target)]
            
            # 更新目标
            target.update(dt, pursuing_usvs)
            self.env.apply_boundary_rules(target)
            
        # 更新USV
        for usv in self.usvs:
            # 获取分配目标
            assigned_target_id = self.usv_assignments.get(id(usv))
            assigned_target=None
            
            # 寻找目标对象
            if assigned_target_id is not None:
                for target in self.targets:
                    if id(target) == assigned_target_id:
                        assigned_target = target
                        break
            
            # 重新分配判断
            if (assigned_target is None or assigned_target.is_captured) and self.enable_reassignment:
                nearest_uncaptured = None
                min_dist=float('inf')
                for target in self.targets:
                    if not target.is_captured:
                        dist = usv.distance_to(target)
                        if dist < min_dist:
                            min_dist = dist
                            nearest_uncaptured = target
                
                if nearest_uncaptured:
                    assigned_target = nearest_uncaptured
                    self.usv_assignments[id(usv)] = id(nearest_uncaptured)
            
            # 获取领导信息
            usv_id = id(usv)
            leadership_info = self.usv_leadership.get(usv_id)
            
            # 更新无人艇
            usv.update(dt, self.usvs, assigned_target, leadership_info)
            self.env.apply_boundary_rules(usv)
            
        # 更新时间
        self.time += dt
        
        # 计算指标
        self._update_metrics()
        
    def _update_metrics(self):
        """更新指标"""
        if not self.targets or not self.usvs:
            return
            
        # 计算距离
        avg_distances = []
        for target in self.targets:
            distances = [usv.distance_to(target) for usv in self.usvs]
            avg_distances.append(np.mean(distances))
        
        self.metrics['avg_distance_to_target'].append(np.mean(avg_distances))
        
        # 计算围捕比例
        surrounded_ratios = []
        for target in self.targets:
            # 8个方向扇区
            sectors_covered = [False] * 8
            
            for usv in self.usvs:
                if usv.distance_to(target) <= usv.percept_radius:
                    # 计算方向角
                    dx = usv.position[0] - target.position[0]
                    dy = usv.position[1] - target.position[1]
                    angle = np.arctan2(dy, dx)
                    
                    # 映射到扇区
                    sector = int(((angle + np.pi) / (2 * np.pi / 8)) % 8)
                    sectors_covered[sector] = True
                    
            surrounded_ratios.append(sum(sectors_covered) / 8)
            
        self.metrics['target_surrounded_ratio'].append(np.mean(surrounded_ratios))
        
        # 计算聚集度
        if len(self.usvs) > 1:
            distances = []
            for i in range(len(self.usvs)):
                for j in range(i+1, len(self.usvs)):
                    distances.append(self.usvs[i].distance_to(self.usvs[j]))
                    
            self.metrics['usv_clustering'].append(np.mean(distances))
        else:
            self.metrics['usv_clustering'].append(0)
            
    def get_stats(self) -> Dict[str, Any]:
        """
        获取当前仿真统计数据
        
        :return: 包含各项统计数据的字典
        """
        stats = {
            'time': self.time,
            'usv_count': len(self.usvs),
            'target_count': len(self.targets),
            'metrics': self.metrics,
        }
        
        # 添加最新指标
        if self.metrics['avg_distance_to_target']:
            stats['current_avg_distance'] = self.metrics['avg_distance_to_target'][-1]
            
        if self.metrics['target_surrounded_ratio']:
            stats['current_surrounded_ratio'] = self.metrics['target_surrounded_ratio'][-1]
            
        # 计算完成情况
        if self.targets and self.metrics['target_surrounded_ratio']:
            surrounded_threshold = 0.8  # 80%的方向被覆盖视为已围捕
            stats['targets_fully_surrounded'] = sum(1 for _ in range(len(self.targets)) 
                                             if self.metrics['target_surrounded_ratio'][-1] >= surrounded_threshold)
            
        return stats 

    def _assign_usvs_to_targets(self):
        """
        智能分配无人艇到目标，确保有效的包围策略
        
        根据当前仿真状态，将无人艇分配给不同目标，并为每个无人艇分配特定的围捕角色。
        该方法会清空当前的角色分配，但保留已有的目标分配以保持稳定性。
        
        :return: None
        """
        # 首先清空当前的角色分配，但保留已有的目标分配用于稳定性
        self.usv_leadership.clear()
        self.encirclement_sectors.clear()
        
        # 获取所有未捕获的目标
        uncaptured_targets = [t for t in self.targets if not t.is_captured]
        if not uncaptured_targets:
            return
            
        # 统计每个目标已分配的无人艇数量
        target_usv_counts = {}
        for target in uncaptured_targets:
            target_id = id(target)
            target_usv_counts[target_id] = 0
        
        for usv_id, target_id in self.usv_assignments.items():
            if target_id in target_usv_counts:
                target_usv_counts[target_id] += 1

        # 打印目标分配情况
        target_usv_counts = {}
        for target in self.targets:
            target_id = id(target)
            target_usv_counts[target_id] = 0
        
        for usv_id, target_id in self.usv_assignments.items():
            if target_id in target_usv_counts:
                target_usv_counts[target_id] += 1
        
        print(f"目标分配情况:")
        for i, target in enumerate(self.targets):
            target_id = id(target)
            count = target_usv_counts.get(target_id, 0)
            print(f"  目标 {i+1}: {count} 艘无人艇")
    
        
        # 确定可用的无人艇（未分配或分配到已捕获目标）和需要重新平衡的无人艇
        available_usvs = []
        assigned_usvs = {}  # 按目标ID分组的已分配无人艇
        
        for target in uncaptured_targets:
            assigned_usvs[id(target)] = []
            
        for usv in self.usvs:
            usv_id = id(usv)
            needs_reassignment = False
            
            # 如果无人艇未被分配
            if usv_id not in self.usv_assignments:
                needs_reassignment = True
            else:
                # 检查分配的目标是否存在且未被捕获
                target_id = self.usv_assignments[usv_id]
                target_exists = False
                for target in uncaptured_targets:
                    if id(target) == target_id:
                        target_exists = True
                        assigned_usvs[target_id].append(usv)
                        break
                
                if not target_exists:
                    needs_reassignment = True
            
            if needs_reassignment:
                available_usvs.append(usv)
        
        # 计算理想的分配情况 - 平均分配无人艇
        total_usvs = len(self.usvs)
        target_count = len(uncaptured_targets)
        ideal_per_target = total_usvs // target_count
        
        # 识别分配不平衡的情况
        targets_needing_usvs = []  # 需要更多无人艇的目标
        targets_with_excess = []   # 有过多无人艇的目标
        
        for target in uncaptured_targets:
            target_id = id(target)
            current_count = len(assigned_usvs[target_id])
            
            if current_count < ideal_per_target - 1:  # 允许±1的波动
                # 需要更多无人艇
                targets_needing_usvs.append((target, ideal_per_target - current_count))
            elif current_count > ideal_per_target + 1:
                # 有过多无人艇
                targets_with_excess.append((target, current_count - ideal_per_target))
        
        # 重平衡：从过剩目标释放无人艇
        newly_available = []
        
        for target, excess in targets_with_excess:
            target_id = id(target)
            # 计算要释放的数量（最多释放过剩数量的75%）
            to_release = min(excess, max(1, int(excess * 0.75)))
            
            # 从这个目标释放无人艇
            usv_to_release = assigned_usvs[target_id][:to_release]
            assigned_usvs[target_id] = assigned_usvs[target_id][to_release:]
            
            # 更新分配
            for usv in usv_to_release:
                usv_id = id(usv)
                if usv_id in self.usv_assignments:
                    del self.usv_assignments[usv_id]
                newly_available.append(usv)
        
        # 添加到可用列表
        available_usvs.extend(newly_available)
        
        # 对目标进行优先级排序（按照已分配的无人艇数量）
        targets_needing_usvs.sort(key=lambda x: x[1], reverse=True)
        
        # 分配可用的无人艇到需要的目标
        for target, needed in targets_needing_usvs:
            if not available_usvs:
                break
                
            target_id = id(target)
            # 计算要分配的数量
            to_assign = min(needed, len(available_usvs))
            
            # 分配无人艇
            for i in range(to_assign):
                if i < len(available_usvs):
                    usv = available_usvs[i]
                    self.usv_assignments[id(usv)] = target_id
            
            # 从可用列表中移除已分配的无人艇
            available_usvs = available_usvs[to_assign:]
        
        # 如果还有剩余的可用无人艇，分配到目标数量最少的目标
        if available_usvs:
            # 按照当前分配数量排序
            sorted_targets = sorted(uncaptured_targets, 
                                   key=lambda t: len(assigned_usvs[id(t)]))
            
            # 轮流分配
            target_index = 0
            for usv in available_usvs:
                target = sorted_targets[target_index]
                self.usv_assignments[id(usv)] = id(target)
                
                # 移到下一个目标
                target_index = (target_index + 1) % len(sorted_targets)
        
        # 最后，为每个目标创建扇区并分配角色
        for target in uncaptured_targets:
            target_id = id(target)
            # 重新计算当前目标的无人艇数量
            target_usv_count = sum(1 for usv_id, tid in self.usv_assignments.items() 
                                if tid == target_id)
            
            # 为目标创建扇区分配
            self._create_encirclement_sectors(target, target_usv_count)
        
        # 为每个目标的无人艇分配包围角色
        self._assign_encirclement_roles()
                
    def _create_encirclement_sectors(self, target, usv_count):
        """
        为目标创建包围扇区
        
        将目标周围360度划分为多个扇区，用于后续分配不同无人艇到不同扇区位置，
        从而形成有效的包围圈。每个扇区包含中心角度、扇区大小和分配状态。
        
        :param target: 目标对象
        :param usv_count: 分配给此目标的无人艇数量
        :return: None
        """
        target_id = id(target)
        
        # 至少需要3个无人艇才能形成有效包围
        if usv_count < 3:
            return
        
        # 创建均匀的扇区
        sectors = []
        for i in range(usv_count):
            # 计算扇区中心角度
            angle = 2 * np.pi * i / usv_count
            # 扇区大小（弧度）
            sector_size = 2 * np.pi / usv_count
            sectors.append({
                'center_angle': angle,
                'size': sector_size,
                'assigned': False  # 标记是否已分配无人艇
            })
        
        self.encirclement_sectors[target_id] = sectors
    
    def _assign_encirclement_roles(self):
        """为各目标分配无人艇的围捕角色"""
        for target in self.targets:
            if target.is_captured:
                continue
                
            target_id = id(target)
            if target_id not in self.encirclement_sectors:
                continue
                
            # 获取分配给该目标的所有无人艇
            assigned_usvs = []
            for usv in self.usvs:
                usv_id = id(usv)
                if self.usv_assignments.get(usv_id) == target_id:
                    assigned_usvs.append(usv)
            
            if not assigned_usvs:
                continue
                
            # 获取该目标的扇区信息
            sectors = self.encirclement_sectors[target_id]
            
            # 按照与目标的角度对无人艇排序
            usv_angles = []
            for usv in assigned_usvs:
                dx = usv.position[0] - target.position[0]
                dy = usv.position[1] - target.position[1]
                angle = (np.arctan2(dy, dx) + 2 * np.pi) % (2 * np.pi)  # 转换到0-2π
                usv_angles.append((usv, angle))
            
            # 为每个扇区分配最接近的无人艇
            for sector in sectors:
                best_usv = None
                min_angle_diff = float('inf')
                sector_angle = sector['center_angle']
                
                for usv, angle in usv_angles:
                    # 计算角度差（考虑循环）
                    angle_diff = min(
                        abs(angle - sector_angle),
                        2 * np.pi - abs(angle - sector_angle)
                    )
                    if angle_diff < min_angle_diff:
                        min_angle_diff = angle_diff
                        best_usv = usv
                
                if best_usv:
                    # 分配角色
                    usv_id = id(best_usv)
                    self.usv_leadership[usv_id] = {
                        'role': 'encirclement',
                        'target_id': target_id,
                        'sector_angle': sector_angle
                    }
                    
                    # 标记为已分配并从可用列表中移除
                    sector['assigned'] = True
                    usv_angles = [(u, a) for u, a in usv_angles if u is not best_usv]
                    
                    # 如果所有无人艇已分配完，结束循环
                    if not usv_angles:
                        break
            
            # 为剩余的无人艇分配备用角色
            for usv, _ in usv_angles:
                usv_id = id(usv)
                # 随机选择一个扇区作为备份
                unassigned_sectors = [s for s in sectors if not s['assigned']]
                if unassigned_sectors:
                    sector = np.random.choice(unassigned_sectors)
                    self.usv_leadership[usv_id] = {
                        'role': 'backup',
                        'target_id': target_id,
                        'sector_angle': sector['center_angle']
                    }
                    sector['assigned'] = True
                else:
                    # 如果所有扇区已分配，则给这个无人艇一个"自由"角色
                    self.usv_leadership[usv_id] = {
                        'role': 'free',
                        'target_id': target_id
                    } 